/*
 * This is written by Lulu Han.
 * E-mail: locomotive_crypto@163.com
 */

#include <iostream>
#include <vector>
#include <exception>
#include "sec_mul.h"
#include "timer.h"
#include "net_client.h"
#include "net_server.h"

void test_mul_for_correctness(long bitlens);

void test_mul_for_performance(long bits, long cnt);

void test_mul_by_network_SA(long bits, long cnt, size_t port, std::string ipaddr);

void test_mul_by_network_SB(size_t port, long vlens);

int main(int argc, char* argv[]){
	std::cout << "----------------------------------------------------------" << std::endl;
	std::cout << "-----------------------<Usage>----------------------------" << std::endl;
	std::cout << "-----1) Run for correctness: ./program--------------------" << std::endl;
	std::cout << "-----2) Run for performance: ./program cnt----------------" << std::endl;
	std::cout << "-----3) Run as SB: ./program port ------------------------" << std::endl;
	std::cout << "-----4) Run as SA: ./program port ipaddr------------------" << std::endl;
	std::cout << "----------------------------------------------------------" << std::endl << std::endl;
	
	if(argc == 1){
		test_mul_for_correctness(200);
	}
	else if(argc == 2){
		for(size_t i = 0; i < 5; ++i)
			test_mul_for_performance(512, std::atoi(argv[1]));
		std::cout << std::endl;
		
		for(size_t i = 0; i < 5; ++i)
			test_mul_for_performance(1024, std::atoi(argv[1]));
		std::cout << std::endl;
		
		for(size_t i = 0; i < 5; ++i)
			test_mul_for_performance(1536, std::atoi(argv[1]));
		std::cout << std::endl;
	}
	else if(argc == 3){
		if(std::string(argv[1]) == std::string("sb") || std::string(argv[1]) == std::string("SB")){
			test_mul_by_network_SB(std::atoi(argv[2]), 10);
		}
			
	}
	else if(argc == 5){
		if(std::string(argv[1]) == std::string("sa") || std::string(argv[1]) == std::string("SA")){
			//test_mul_by_network_SA(512, std::atoi(argv[4]), std::atoi(argv[2]), std::string(argv[3]));
			
			//test_mul_by_network_SA(1024, std::atoi(argv[4]), std::atoi(argv[2]), std::string(argv[3]));
			
			test_mul_by_network_SA(1536, std::atoi(argv[4]), std::atoi(argv[2]), std::string(argv[3]));
		}
	}
	return 0;
}

void test_mul_for_correctness(long bitlens){
	PrivateKey prikey;
	PublicKey pubkey;
	
	Paillier::GenKeyPair(prikey, pubkey, bitlens);
	
	Decryptor dec(prikey);
	Encryptor enc(pubkey);
	
	NTL::ZZ a(12345), b(67890);
	
	std::cout << "a = " << a << std::endl;
	std::cout << "b = " << b << std::endl;
	std::cout << "a*b = " << a*b << std::endl;
	
	NTL::ZZ ca, cb;
	enc.Encrypt(a, ca); // ca = Enc(a)
	enc.Encrypt(b, cb); // cb = Enc(b)
	
	SecMul::CloudA CA(prikey);
	SecMul::CloudB CB(pubkey);
	
	NTL::ZZ nca, ncb;
	CB.gen_two_ciphers(ca, cb, nca, ncb);
	
	NTL::ZZ cmul;
	CA.dec_reenc(nca, ncb, cmul);
	
	NTL::ZZ cab;
	CB.gen_mul(cmul, cab);
	std::cout << "Enc(a + b)=" << cab << std::endl;
	
	NTL::ZZ ab;
	dec.Decrypt(cab, ab);
	std::cout << "ab=" << ab << std::endl;
	
	if(a*b != ab){
		throw std::runtime_error("Failed to test SecMul ...\n");
	}
}

void test_mul_for_performance(long bits, long cnt){
	std::cout << "The key length for this test is: " << 2*bits << std::endl;
	 
	std::vector<Paillier::Encryptor> vencryptor;
	std::vector<Paillier::Decryptor> vdecryptor;
	
	std::vector<SecMul::CloudA> vCA;
	std::vector<SecMul::CloudB> vCB;
	
	for(long i = 0; i < cnt; ++i){
		Paillier::PrivateKey prikey;
		Paillier::PublicKey pubkey;
	
		// Generate random key pair
		Paillier::GenKeyPair(prikey, pubkey, bits);
	
		// Define Encryptor and Decryptor objects
		Paillier::Encryptor encryptor(pubkey);
		Paillier::Decryptor decryptor(prikey);
		
		vencryptor.push_back(encryptor);
		vdecryptor.push_back(decryptor);
		
		SecMul::CloudA CA(prikey);
		SecMul::CloudB CB(pubkey);
		
		vCA.push_back(CA);
		vCB.push_back(CB);
	}
	
	std::vector<NTL::ZZ> va, vb;
	
	NTL::ZZ r;
	for(long i = 0; i < cnt; ++i){
		r = NTL::RandomBnd(vencryptor[i].GetN());
		va.push_back(r);
		
		r = NTL::RandomBnd(vencryptor[i].GetN());
		vb.push_back(r);
	}
	
	std::vector<NTL::ZZ> vca, vcb;
	NTL::ZZ c;
	for(long i = 0; i < cnt; ++i){
		vencryptor[i].Encrypt(va[i], c);
		vca.push_back(c);
		
		vencryptor[i].Encrypt(vb[i], c);
		vcb.push_back(c);
	}
	
	time_point<high_resolution_clock> start;	
	start = now_time();
	for(long i = 0; i < cnt; ++i){
		NTL::ZZ nca, ncb;
		vCB[i].gen_two_ciphers(vca[i], vcb[i], nca, ncb);
		
		NTL::ZZ cmul;
		vCA[i].dec_reenc(nca, ncb, cmul);
		
		NTL::ZZ cab;
		vCB[i].gen_mul(cmul, cab);
	}
	double tm = time_from(start);
	std::cout << "The time for "<< cnt << " SecMul is : " << tm << std::endl;
}

void test_mul_by_network_SA(long bits, long cnt, size_t port, std::string ipaddr){
	std::cout << "The key length for this test is: " << 2*bits << std::endl;
	
	NetClient client(ipaddr, port);
	client.connect();
	
	client.send(cnt);
	client.send(bits);
	
	std::vector<Paillier::Decryptor> vdecryptor;
	
	std::vector<SecMul::CloudA> vCA;
	
	for(long i = 0; i < cnt; ++i){
		Paillier::PrivateKey prikey;
		Paillier::PublicKey pubkey;
	
		// Generate random key pair
		Paillier::GenKeyPair(prikey, pubkey, bits);
		
		client.send(pubkey.GetN());
		
		// Define Encryptor and Decryptor objects
		Paillier::Decryptor decryptor(prikey);
		
		vdecryptor.push_back(decryptor);
		
		SecMul::CloudA CA(prikey);
		
		vCA.push_back(CA);
	}
	
	std::cout << "Please count the network traffics from here..." << std::endl;
	getchar();
	for(long i = 0; i < cnt; ++i){
		NTL::ZZ nca, ncb;
		client.receive(nca);
		client.receive(ncb);
		
		NTL::ZZ cmul;
		vCA[i].dec_reenc(nca, ncb, cmul);
		client.send(cmul);
		
		/*
		NTL::ZZ cab;
		client.receive(cab);
		
		NTL::ZZ pab;
		vdecryptor[i].Decrypt(cab, pab);
		std::cout << "pab = " << pab << std::endl;
		*/
	}
	
	client.close();
}

void test_mul_by_network_SB(size_t port, long vlens){
	NetServer server(port); 
	server.accept();
	
	long cnt;
	server.receive(cnt);
	
	long bits;
	server.receive(bits);
	
	std::cout << "The key length for this test is: " << 2*bits << std::endl;
	
	std::vector<Paillier::Encryptor> vencryptor;
	std::vector<SecMul::CloudB> vCB;
	
	NTL::ZZ N;
	for(long i = 0; i < cnt; ++i){
		server.receive(N);
		
		Paillier::PublicKey pubkey(N);
		SecMul::CloudB CB(pubkey);
		
		Paillier::Encryptor enc(pubkey);
		vencryptor.push_back(enc);
		
		vCB.push_back(CB);
	}
	
	std::vector<NTL::ZZ> vx, vy;
	
	NTL::ZZ r;
	for(long i = 0; i < cnt; ++i){
		r = NTL::RandomLen_ZZ(30);
		vx.push_back(r);
		
		r = NTL::RandomLen_ZZ(30);
		vy.push_back(r);
	}
	
	std::vector<NTL::ZZ> vcx, vcy;
	NTL::ZZ c;
	for(long i = 0; i < cnt; ++i){
		vencryptor[i].Encrypt(vx[i], c);
		vcx.push_back(c);
		
		vencryptor[i].Encrypt(vy[i], c);
		vcy.push_back(c);
	}
	
	std::cout << "Please count the network traffics from here..." << std::endl;
	getchar();
	for(long i = 0; i < cnt; ++i){
		NTL::ZZ nca, ncb;
		vCB[i].gen_two_ciphers(vcx[i], vcy[i], nca, ncb);
		server.send(nca);
		server.send(ncb);
		
		NTL::ZZ cmul;
		server.receive(cmul);
		
		NTL::ZZ cab;
		vCB[i].gen_mul(cmul, cab);
		//server.send(cab);
		
		//std::cout << "x*y = " << vx[i] * vy[i] << std::endl;
	}
	
	server.close();
}
